#include "windivert_driver.h"
#include "lwip/ip.h"
#include "crc32.h"

static const uint32_t APP_NAMES_SIZE = 4096;
using namespace std::chrono_literals;

void WindivertDriver::_runWrite()
{
    std::unique_ptr<WINDIVERT_ADDRESS[]> addrs(new WINDIVERT_ADDRESS[100]);
    std::memset(addrs.get(),0,sizeof(WINDIVERT_ADDRESS)*100);
    for (int i = 0;i < 100;++i) {
        addrs[i].Network.IfIdx = 6;
        addrs[i].Network.SubIfIdx = 0;
        addrs[i].Layer = WINDIVERT_LAYER_NETWORK;
        addrs[i].Outbound = 0;
        addrs[i].TCPChecksum = 1;
        addrs[i].UDPChecksum = 1;
        addrs[i].IPChecksum = 1;
        addrs[i].Impostor = 1;
        addrs[i].IPv6 = 0;
    }

    for (;;) {
        if (is_wstop_.test()) {
            is_wstop_.clear();
            is_wstop_.notify_all();
            return;
        }

        if (queue_inject_.empty()) {
            queue_inject_nonempty_.clear();
            queue_inject_nonempty_.wait(false,std::memory_order_relaxed);
            continue;
        }

        uint8_t* data = (uint8_t*)inject_data_.get();
        uint32_t data_size = 0;
        uint32_t addrs_index = 0;
        for (uint16_t i = 0; i < 50; ++i) {
            NetPacket** it = queue_inject_.front();
            if (it == nullptr) {
                Sleep(0);
                continue;
            }

            NetPacket* buffer = *it;
            queue_inject_.pop();
            uint8_t* buffer_data = buffer->data;

            uint32_t key = 0;
            uint8_t ip_version = IP_HDR_GET_VERSION(buffer_data);
            if (ip_version == 4) {
                key = *((uint32_t*)(buffer_data+12));
            } else if (ip_version == 6) {
                addrs[addrs_index].IPv6 = 1;
                key = calc_crc32(0xffffffff,buffer_data+8,16);
            }
            {
                auto ifindex = map_ifindex_.find(key);
                if (ifindex != map_ifindex_.end()) {
                    addrs[addrs_index].Network.IfIdx = ifindex->second;
                } else {
                    std::cerr << "map_ifindex find error!\n";
                }
            }


            std::memcpy(data,buffer_data,buffer->data_len);
            data += buffer->data_len;
            data_size += buffer->data_len;
            ++addrs_index;

            NetPacketPool::deletePacket(buffer);
        }

        uint32_t send_len = 0;
        if (!WinDivertSendEx(w_handle_, inject_data_.get(), data_size, &send_len, 0,
                             addrs.get(), addrs_index*sizeof(WINDIVERT_ADDRESS), NULL) ||
            send_len != data_size) {
            std::cerr << "WinDivertSendEx Fail\n";
        }
    }
}

void WindivertDriver::_runRead()
{

    WINDIVERT_ADDRESS windivert_addr[10];
    uint32_t recv_size = 0;
    for (;;) {
        if (is_rstop_.test()) {
            is_rstop_.clear();
            is_rstop_.notify_all();
            return;
        }
        uint32_t addr_len = 10*sizeof(WINDIVERT_ADDRESS);
        memset(windivert_addr, 0, addr_len);
#if 0
        WINDIVERT_ADDRESS windivert_addr1;
        if (!WinDivertRecv(w_handle_, recv_data_.get(), WINDIVERT_MTU_MAX, &recv_size, &windivert_addr1)) {
            std::cerr << "failed to read packet " << GetLastError() << "\n";
            break;
        }
        uint8_t* data = (uint8_t*)recv_data_.get();
        uint32_t key = 0;
        uint8_t ip_version = IP_HDR_GET_VERSION(data);
        if (ip_version == 4) {
            key = *((uint32_t*)(data+16));
        } else if (ip_version == 6) {
            key = calc_crc32(0xffffffff,data+24,16);
        }
        WinDivertHelperCalcChecksums(data, recv_size, &windivert_addr1, 0);
        map_ifindex_.insert({key, windivert_addr1.Network.IfIdx});

        std::shared_ptr<NetPacket> buf = _NetPacketPool->getSharedPacket(recv_size);
        memcpy(buf->data, recv_data_.get(), recv_size);
        if (cb_out_data_) cb_out_data_(buf,recv_size);
#else
        if (!WinDivertRecvEx(w_handle_, recv_data_.get(), 10*1500, &recv_size, 0, windivert_addr, &addr_len, nullptr)) {
            std::cerr << "failed to read packet " << GetLastError() << "\n";
            break;
        }
        
        uint8_t* data = (uint8_t*)recv_data_.get();
        for (int i = 0;i < addr_len/sizeof(WINDIVERT_ADDRESS);++i) {
            if (windivert_addr[i].Event != WINDIVERT_LAYER_NETWORK) {
                std::cerr << "event != WINDIVERT_LAYER_NETWORK\n";
                continue;
            }
            if (windivert_addr[i].Network.Protocol == WINDIVERT_IP_PROTOCOL_UDP) {
                std::cerr << "Protocol == WINDIVERT_IP_PROTOCOL_UDP\n";
                continue;
            }

            uint16_t packet_len = 0;
            if (getPacketLen(data,packet_len)) {
                WinDivertHelperCalcChecksums(data, packet_len, &windivert_addr[i], 0);

                uint32_t key = 0;
                uint8_t ip_version = IP_HDR_GET_VERSION(data);
                if (ip_version == 4) {
                    key = *((uint32_t*)(data+16));
                } else if (ip_version == 6) {
                    key = calc_crc32(0xffffffff,data+24,16);
                }
                map_ifindex_.insert({key, windivert_addr[i].Network.IfIdx});

                std::shared_ptr<NetPacket> buf = _NetPacketPool->getSharedPacket(packet_len);
                memcpy(buf->data, data, packet_len);
                if (cb_out_data_) cb_out_data_(buf);
                data += packet_len;
            } else {
                std::cerr << "getPacketLen Error\n";
            }
        }
#endif
    }
}

void WindivertDriver::doWrite(NetPacket *buffer)
{
    queue_inject_.emplace(buffer);
    if (!queue_inject_nonempty_.test()) {
        queue_inject_nonempty_.test_and_set();
        queue_inject_nonempty_.notify_all();
    }
#if 0
    WINDIVERT_ADDRESS addr;
    addr.Network.IfIdx = 6;
    addr.Network.SubIfIdx = 0;
    addr.Outbound = 0;
    addr.TCPChecksum = 1;
    addr.UDPChecksum = 1;
    addr.IPChecksum = 1;
    addr.Impostor = 1;
    addr.IPv6 = 0;

    if (IP_HDR_GET_VERSION(buffer->data) == 6) {
        addr.IPv6 = 1;
    }

    WinDivertSend(w_handle_, buffer->data, len, NULL, &addr);
#endif
}

void WindivertDriver::stop(bool wait)
{
    is_rstop_.test_and_set(std::memory_order_acquire);
    if (wait) is_rstop_.wait(true,std::memory_order_relaxed);
    is_wstop_.test_and_set(std::memory_order_acquire);
    if (wait) is_wstop_.wait(true,std::memory_order_relaxed);
}

bool WindivertDriver::getPacketLen(uint8_t *packet, uint16_t &packet_len)
{
    PWINDIVERT_IPHDR ip_header = (PWINDIVERT_IPHDR)packet;
    PWINDIVERT_IPV6HDR ipv6_header = NULL;
    if (ip_header->Version == 4) {
        packet_len = (UINT)lwip_ntohs(ip_header->Length);
        return true;
    } else if (ip_header->Version == 6) {
        ipv6_header = (PWINDIVERT_IPV6HDR)packet;
        packet_len = (UINT)lwip_ntohs(ipv6_header->Length) +
                    sizeof(WINDIVERT_IPV6HDR);
        return true;
    }
    return false;
}

WindivertDriver::WindivertDriver(const std::vector<std::string> &app_names):
    recv_data_(_aligned_malloc(15008, 16),[](void* p) {_aligned_free(p);}),
    inject_data_(_aligned_malloc(1024*1024, 16),[](void* p) {_aligned_free(p);}),
    app_names_((wchar_t*)_aligned_malloc(APP_NAMES_SIZE, 16),[](wchar_t* p) {_aligned_free(p);}),
    queue_inject_(2048)
{
    auto apps_ptr = app_names_.get();
    memset(apps_ptr,0, APP_NAMES_SIZE);
    uint32_t index = 0;
    for (const std::string& app_name : app_names) {
        int char_size = MultiByteToWideChar(CP_UTF8,0,app_name.data(), app_name.length(),apps_ptr+index, APP_NAMES_SIZE -index);
        index += char_size;
        ++index;
        if (index >= APP_NAMES_SIZE) break;
    }

    w_handle_ = WinDivertOpen((const char*)apps_ptr, WINDIVERT_LAYER_NETWORK, 776, 0);
    if (w_handle_ == INVALID_HANDLE_VALUE) {
        SERROR("WinDivertOpen error:INVALID_HANDLE_VALUE");
        stop(false);
        return;
    }
    WinDivertSetParam(w_handle_, WINDIVERT_PARAM_QUEUE_TIME, WINDIVERT_PARAM_QUEUE_TIME_MAX);
    WinDivertSetParam(w_handle_, WINDIVERT_PARAM_QUEUE_LENGTH, WINDIVERT_PARAM_QUEUE_LENGTH_MAX);
    WinDivertSetParam(w_handle_, WINDIVERT_PARAM_QUEUE_SIZE, WINDIVERT_PARAM_QUEUE_SIZE_MAX);
}

WindivertDriver::~WindivertDriver()
{
    stop(true);
    WinDivertClose(w_handle_);

    for (auto& thd : workers_) {
        thd.join();
    }
}

bool WindivertDriver::run(cb_outbound_data out)
{
    if (w_handle_ == INVALID_HANDLE_VALUE) return false;

    cb_out_data_ = out;
    workers_.emplace_back(std::bind(&WindivertDriver::_runRead,this));
    workers_.emplace_back(std::bind(&WindivertDriver::_runWrite,this));
    return true;
}
